from __future__ import with_statement

import sys
import math
import os.path
import logging
import glob
import optparse
import re

import common
import refolding_common

logger = logging.getLogger("refolding_add_wait_info")
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.INFO)

def data_filename(filename):
    return filename.rstrip("*").rstrip("_break")

def read_waiting_time(jp_path):
    with open(jp_path, "rt") as g:
        line = list(g)[0]
    fields = line.split(", ")
    t_s = (int(fields[6]) - int(fields[5])) / 5000.0
    return t_s

def read_hover_nm(info_path):
    with open(info_path, "rt") as g:
        line = list(g)[0]
    return float(line)

spring_constant_pattern = re.compile(
    r"^\\Spring Constant: (?P<k>(?:(?:[0-9]+\.?[0-9]*)|(?:\.[0-9]+))(?:[eE](\+|-)?[0-9]+)?)\s*$")

def extract_spring_constant(data_file_path):
    k = None
    with open(data_file_path, "rt") as f:
        for line in f:
            m = spring_constant_pattern.match(line)
            if m:
                k = refolding_common.float_type(m.group("k"))
                break  # optimisation
    common.check_True(k)
    return k

def extend_to_full(root_dir, name):
    jp_dir = os.path.join(root_dir, "refolding_jp_%s" % name)
    info_dir = os.path.join(root_dir, "refolding_info_%s" % name)
    data_dir = os.path.join(root_dir, "refolding_selected_%s" % name)
    full = []
    entries = refolding_common.read_entries(os.path.join(root_dir, name + ".txt"))
    k = None
    for entry in entries:
        common.check_True(isinstance(entry, tuple))
        data_filename = refolding_common.data_filename(entry[-1])
        jp_path = os.path.join(jp_dir, data_filename + ".junctionPoints")
        t_s = read_waiting_time(jp_path)
        info_path = os.path.join(info_dir, data_filename + ".info")
        hover_nm = read_hover_nm(info_path)
        k = extract_spring_constant(os.path.join(data_dir, data_filename))
        common.check_False(k is None)
        full_entry = (k, hover_nm, t_s * 1e3) + entry
        full.append(full_entry)
    with open(os.path.join(root_dir, name + ".full"), "wt") as f:
        for full_entry in full:
            f.write("%s\n" % (", ".join(map(repr, full_entry))))

def main(params):
    parser = optparse.OptionParser()
    parser.add_option(
        "--root_dirs", help="dirs having under them .txt files and junction points directories") 
    options, args = parser.parse_args(params)
    if args:
        parser.print_help()
        return -1
    print(options)

    root_dirs = options.root_dirs
    for root_dir in root_dirs.split(":"):
        for filepath in glob.iglob(os.path.join(root_dir, "*.txt")):
            name = os.path.basename(os.path.splitext(filepath)[0])
            extend_to_full(root_dir, name)

if __name__ == "__main__":
    main(sys.argv[1:])
